"""
茅台股票预测系统 - Streamlit主应用
🍶 基于机器学习的茅台股票价格预测系统
"""

import streamlit as st
import sys
import os

# 添加src目录到路径
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))

from data_loader import (
    get_stock_data, add_technical_indicators, get_stock_info,
    get_stock_options, get_industry_options, validate_stock_code,
    get_stock_suggestions
)
from predictor import StockPredictor
from visualizer import (
    plot_stock_price, plot_candlestick, plot_technical_indicators,
    plot_prediction_results, plot_model_performance, display_metrics
)

# 页面配置
st.set_page_config(
    page_title="A股股票预测系统",
    page_icon="📈",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义CSS
st.markdown("""
<style>
    .main-header {
        text-align: center;
        padding: 1rem;
        background: linear-gradient(90deg, #ff6b6b, #4ecdc4);
        border-radius: 10px;
        margin-bottom: 2rem;
    }
    .metric-card {
        background: #f0f2f6;
        padding: 1rem;
        border-radius: 10px;
        border-left: 4px solid #ff6b6b;
    }
    .sidebar .sidebar-content {
        background: linear-gradient(180deg, #667eea 0%, #764ba2 100%);
    }
</style>
""", unsafe_allow_html=True)

# 主标题
st.markdown("""
<div class="main-header">
    <h1>📈 A股股票预测系统</h1>
    <p>基于机器学习的股价预测 | 数据驱动投资决策</p>
</div>
""", unsafe_allow_html=True)

# 侧边栏配置
st.sidebar.title("⚙️ 系统配置")

# 股票选择
st.sidebar.subheader("🏢 股票选择")

# 快速代码输入
st.sidebar.markdown("#### 🚀 快速输入")
quick_code_input = st.sidebar.text_input(
    "直接输入股票代码",
    placeholder="如: 600519, 000001, 300750",
    help="输入6位数字A股代码，回车确认"
)

# 实时验证和显示
if quick_code_input:
    is_valid, stock_info = validate_stock_code(quick_code_input)
    if is_valid:
        st.sidebar.success(f"✅ {stock_info.get('name', '股票')} ({stock_info['code']})")
        # 使用快速输入的股票
        selected_stock_code = stock_info['code']
        selected_stock_name = stock_info.get('name', f'股票{stock_info["code"]}')
        selection_mode_used = "快速输入"
    else:
        st.sidebar.error("⚠️ 无效的股票代码")
        # 显示建议
        suggestions = get_stock_suggestions(quick_code_input)
        if suggestions:
            st.sidebar.info(f"💡 您是否想要: {', '.join(suggestions[:3])}")
else:
    selection_mode_used = None

# 分割线
st.sidebar.markdown("---")
st.sidebar.markdown("#### 📊 或者选择以下方式")

# 获取股票选项
stock_names, stock_codes = get_stock_options()
industry_stocks = get_industry_options()

# 选择模式（只在没有快速输入时显示）
if not quick_code_input or not validate_stock_code(quick_code_input)[0]:
    selection_mode = st.sidebar.radio(
        "选择模式",
        ["🔥 热门股票", "🏢 行业分类", "✍️ 手动输入"],
        disabled=bool(quick_code_input and validate_stock_code(quick_code_input)[0])
    )
    
    if selection_mode == "🔥 热门股票":
        selected_stock_name = st.sidebar.selectbox(
            "选择股票",
            stock_names,
            index=0  # 默认选择茅台
        )
        selected_stock_code = stock_codes[selected_stock_name]
        selection_mode_used = "热门股票"
        
    elif selection_mode == "🏢 行业分类":
        selected_industry = st.sidebar.selectbox(
            "选择行业",
            list(industry_stocks.keys())
        )
        
        industry_stock_names = list(industry_stocks[selected_industry].keys())
        selected_stock_name = st.sidebar.selectbox(
            "选择股票",
            industry_stock_names
        )
        selected_stock_code = industry_stocks[selected_industry][selected_stock_name]
        selection_mode_used = f"行业分类 - {selected_industry}"
        
    else:  # 手动输入
        selected_stock_name = st.sidebar.text_input(
            "股票名称",
            value="贵州茅台"
        )
        
        manual_stock_code = st.sidebar.text_input(
            "股票代码",
            value="600519",
            help="请输入6位数字的A股代码，如600519"
        )
        
        # 验证代码格式
        is_valid_manual, _ = validate_stock_code(manual_stock_code)
        if is_valid_manual:
            selected_stock_code = manual_stock_code
            st.sidebar.success("✅ 代码格式正确")
        else:
            st.sidebar.error("⚠️ 股票代码格式不正确！")
            selected_stock_code = "600519"
        
        selection_mode_used = "手动输入"
else:
    # 已经通过快速输入选择了股票
    pass

# 显示当前选择
if 'selected_stock_code' in locals() and 'selected_stock_name' in locals():
    st.sidebar.info(f"🎯 当前选择: {selected_stock_name} ({selected_stock_code})")
    if selection_mode_used:
        st.sidebar.caption(f"📋 选择方式: {selection_mode_used}")
else:
    # 默认选择茅台
    selected_stock_name = "贵州茅台"
    selected_stock_code = "600519"
    st.sidebar.warning("⚠️ 使用默认选择: 贵州茅台(600519)")

# 数据获取配置
st.sidebar.subheader("📊 数据配置")
data_days = st.sidebar.slider("历史数据天数", 365, 1095, 730, 30)
auto_refresh = st.sidebar.checkbox("自动刷新数据", value=False)

# 预测配置
st.sidebar.subheader("🔮 预测配置")
predict_days = st.sidebar.slider("预测天数", 10, 30, 20, 1)
sequence_length = st.sidebar.slider("时序长度", 5, 15, 10, 1)
test_ratio = st.sidebar.slider("测试集比例", 0.1, 0.3, 0.2, 0.05)

# 显示配置
st.sidebar.subheader("📈 显示配置")
show_candlestick = st.sidebar.checkbox("显示K线图", value=True)
show_technical = st.sidebar.checkbox("显示技术指标", value=True)
show_performance = st.sidebar.checkbox("显示模型性能", value=True)

# 初始化session state
if 'df' not in st.session_state:
    st.session_state.df = None
if 'predictor' not in st.session_state:
    st.session_state.predictor = None
if 'model_trained' not in st.session_state:
    st.session_state.model_trained = False
if 'prediction_results' not in st.session_state:
    st.session_state.prediction_results = None
if 'current_stock' not in st.session_state:
    st.session_state.current_stock = None

# 数据加载
@st.cache_data(ttl=3600)
def load_data(stock_code, stock_name, days):
    """加载股票数据"""
    df, data_type = get_stock_data(stock_code, stock_name, days)
    df = add_technical_indicators(df)
    return df, data_type

# 主要功能选项卡
tab1, tab2, tab3, tab4 = st.tabs(["📊 数据总览", "📈 技术分析", "🔮 智能预测", "📋 模型评估"])

# 数据加载按钮
col1, col2, col3 = st.columns([1, 2, 1])
with col2:
    if st.button(f"🔄 加载 {selected_stock_name} 数据", type="primary", use_container_width=True):
        # 检查是否更换了股票
        current_stock_key = f"{selected_stock_code}_{selected_stock_name}"
        
        if st.session_state.current_stock != current_stock_key:
            # 更换股票时清空之前的模型和预测
            st.session_state.model_trained = False
            st.session_state.prediction_results = None
            st.session_state.predictor = None
            st.session_state.current_stock = current_stock_key
        
        with st.spinner(f"正在获取{selected_stock_name}股票数据..."):
            st.session_state.df, data_type = load_data(selected_stock_code, selected_stock_name, data_days)
            st.success(f"✅ 成功获取{selected_stock_name}{data_type}！数据量：{len(st.session_state.df)}条")

# 检查是否有数据
if st.session_state.df is None:
    st.warning("⚠️ 请先选择股票并点击加载数据按钮")
    st.info("💡 您可以在左侧侧边栏中选择不同的股票进行分析")
    st.stop()

df = st.session_state.df
stock_info = get_stock_info(df, selected_stock_name)

# Tab 1: 数据总览
with tab1:
    st.subheader(f"📊 {selected_stock_name}({selected_stock_code}) 数据总览")
    
    # 基本信息卡片
    col1, col2, col3, col4 = st.columns(4)
    
    with col1:
        st.metric(
            label="💰 当前价格",
            value=f"¥{stock_info['current_price']:.2f}",
            delta=f"{stock_info['change']:+.2f} ({stock_info['change_pct']:+.2f}%)"
        )
    
    with col2:
        st.metric(
            label="📈 52周最高",
            value=f"¥{stock_info['high_52w']:.2f}"
        )
    
    with col3:
        st.metric(
            label="📉 52周最低",
            value=f"¥{stock_info['low_52w']:.2f}"
        )
    
    with col4:
        st.metric(
            label="📊 平均成交量",
            value=f"{stock_info['volume_avg']:.0f}"
        )
    
    # 技术指标当前值
    col1, col2, col3 = st.columns(3)
    with col1:
        if 'ma5' in stock_info:
            st.metric("MA5", f"¥{stock_info['ma5']:.2f}")
    with col2:
        if 'ma20' in stock_info:
            st.metric("MA20", f"¥{stock_info['ma20']:.2f}")
    with col3:
        if 'rsi' in stock_info:
            rsi_value = stock_info['rsi']
            rsi_status = "超买" if rsi_value > 70 else "超卖" if rsi_value < 30 else "正常"
            st.metric("RSI", f"{rsi_value:.1f}", delta=rsi_status)
    
    # 数据信息
    st.info(f"📅 数据范围: {stock_info['data_start']} 至 {stock_info['data_end']} | 总记录数: {stock_info['total_records']}")
    
    # 价格走势图
    st.subheader("💹 价格走势")
    price_chart = plot_stock_price(df, f"{selected_stock_name}股票价格走势")
    st.plotly_chart(price_chart, use_container_width=True)
    
    # K线图（可选）
    if show_candlestick:
        st.subheader("🕯️ K线图")
        candlestick_chart = plot_candlestick(df, f"{selected_stock_name}股票K线图")
        st.plotly_chart(candlestick_chart, use_container_width=True)

# Tab 2: 技术分析
with tab2:
    st.subheader("📈 技术指标分析")
    
    if show_technical:
        # 技术指标图表
        tech_chart = plot_technical_indicators(df)
        st.plotly_chart(tech_chart, use_container_width=True)
        
        # 技术分析建议
        st.subheader("💡 技术分析建议")
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.markdown("#### 🎯 移动平均线分析")
            current_price = stock_info['current_price']
            ma5 = stock_info.get('ma5', 0)
            ma20 = stock_info.get('ma20', 0)
            
            if current_price > ma5 > ma20:
                st.success("📈 多头排列：价格 > MA5 > MA20，趋势向上")
            elif current_price < ma5 < ma20:
                st.error("📉 空头排列：价格 < MA5 < MA20，趋势向下")
            else:
                st.warning("🔄 震荡格局：均线交错，趋势不明")
        
        with col2:
            st.markdown("#### 📊 RSI 强弱分析")
            rsi = stock_info.get('rsi', 50)
            
            if rsi > 70:
                st.error("⚠️ RSI > 70：超买状态，谨防回调")
            elif rsi < 30:
                st.success("💰 RSI < 30：超卖状态，关注反弹")
            else:
                st.info(f"✅ RSI = {rsi:.1f}：正常区间，可持续观察")
    
    # 数据表格
    st.subheader("📋 最新数据")
    st.dataframe(
        df.tail(10)[['open', 'high', 'low', 'close', 'volume', 'ma5', 'ma20', 'rsi']].round(2),
        use_container_width=True
    )

# Tab 3: 智能预测
with tab3:
    st.subheader("🔮 AI智能预测")
    
    # 模型训练
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        if st.button("🧠 训练预测模型", type="primary", use_container_width=True):
            with st.spinner("正在训练机器学习模型..."):
                try:
                    # 初始化预测器
                    st.session_state.predictor = StockPredictor(sequence_length=sequence_length)
                    
                    # 训练模型
                    train_results = st.session_state.predictor.train(df, test_ratio=test_ratio)
                    st.session_state.model_trained = True
                    
                    # 显示训练结果
                    st.success("✅ 模型训练完成！")
                    
                    # 显示评估指标
                    display_metrics(train_results['test_metrics'], "🎯 模型测试性能")
                    
                except Exception as e:
                    st.error(f"❌ 模型训练失败: {str(e)}")
    
    # 预测功能
    if st.session_state.model_trained and st.session_state.predictor:
        
        col1, col2, col3 = st.columns([1, 2, 1])
        with col2:
            if st.button(f"🚀 预测未来{predict_days}天", type="secondary", use_container_width=True):
                with st.spinner(f"正在预测未来{predict_days}天股价..."):
                    try:
                        # 执行预测
                        prediction_results = st.session_state.predictor.predict_future(df, days=predict_days)
                        st.session_state.prediction_results = prediction_results
                        
                        st.success(f"✅ 预测完成！未来{predict_days}天股价预测已生成")
                        
                    except Exception as e:
                        st.error(f"❌ 预测失败: {str(e)}")
        
        # 显示预测结果
        if st.session_state.prediction_results:
            results = st.session_state.prediction_results
            
            # 预测摘要
            st.subheader("📊 预测结果摘要")
            
            col1, col2, col3, col4 = st.columns(4)
            
            with col1:
                st.metric(
                    "当前价格",
                    f"¥{results['stats']['current_price']:.2f}"
                )
            
            with col2:
                st.metric(
                    f"第{predict_days}天预测",
                    f"¥{results['stats']['max_predicted']:.2f}",
                    delta=f"{results['stats']['total_change']:+.2f}"
                )
            
            with col3:
                st.metric(
                    "预期涨跌幅",
                    f"{results['stats']['total_change_pct']:+.2f}%"
                )
            
            with col4:
                st.metric(
                    "投资建议",
                    results['suggestion']
                )
            
            # 预测图表
            st.subheader("📈 预测走势图")
            prediction_chart = plot_prediction_results(df, results, f"{selected_stock_name}股票{predict_days}天预测")
            st.plotly_chart(prediction_chart, use_container_width=True)
            
            # 详细预测数据
            st.subheader("📋 详细预测数据")
            
            # 格式化预测数据
            display_df = results['predictions'].copy()
            display_df['date'] = display_df['date'].dt.strftime('%Y-%m-%d')
            display_df['predicted_price'] = display_df['predicted_price'].round(2)
            display_df['change_amount'] = display_df['change_amount'].round(2)
            display_df['change_percent'] = display_df['change_percent'].round(2)
            
            st.dataframe(
                display_df.rename(columns={
                    'date': '日期',
                    'predicted_price': '预测价格(¥)',
                    'change_amount': '变化金额(¥)',
                    'change_percent': '变化比例(%)'
                }),
                use_container_width=True
            )
            
            # 下载预测结果
            csv = results['predictions'].to_csv(index=False)
            st.download_button(
                label="📥 下载预测结果 (CSV)",
                data=csv,
                file_name=f"{selected_stock_code}_{selected_stock_name}_prediction_{predict_days}days.csv",
                mime="text/csv"
            )
    
    else:
        st.info("💡 请先训练模型，然后进行股价预测")

# Tab 4: 模型评估
with tab4:
    st.subheader("📋 模型性能评估")
    
    if st.session_state.model_trained and st.session_state.predictor:
        # 模型信息
        model_info = st.session_state.predictor.get_model_info()
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.markdown("#### 🤖 模型信息")
            st.json(model_info)
        
        with col2:
            st.markdown("#### ⚙️ 训练参数")
            st.write(f"- 时序长度: {sequence_length}")
            st.write(f"- 测试集比例: {test_ratio:.1%}")
            st.write(f"- 特征数量: {len(model_info['features'])}")
            st.write(f"- 数据总量: {len(df)}")
        
        # 特征重要性（简化显示）
        st.markdown("#### 📊 使用特征")
        features_df = pd.DataFrame({
            '特征名称': model_info['features'],
            '特征说明': [
                '开盘价', '最高价', '最低价', '成交量',
                '5日均线', '20日均线', 'RSI指标', '前收盘价'
            ]
        })
        st.dataframe(features_df, use_container_width=True)
        
    else:
        st.info("💡 请先训练模型以查看评估结果")

# 页脚
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #666; padding: 1rem;'>
    <p>⚠️ <strong>风险提示</strong>: 本系统仅供学习和参考，不构成投资建议。股市有风险，投资需谨慎！</p>
    <p>🛠️ 技术栈: Python + Streamlit + AkShare + scikit-learn + Plotly</p>
    <p>🏢 支持股票: 热门A股、行业分类、自定义输入</p>
    <p>📖 项目地址: <a href="https://gitee.com/hnnsboy/simple-stock-prediction" target="_blank">Gitee</a></p>
</div>
""", unsafe_allow_html=True)

# 实时状态显示
with st.sidebar:
    st.markdown("---")
    st.subheader("📊 系统状态")
    
    if st.session_state.df is not None:
        st.success(f"✅ 数据已加载 ({len(st.session_state.df)}条)")
    else:
        st.error("❌ 数据未加载")
    
    if st.session_state.model_trained:
        st.success("✅ 模型已训练")
    else:
        st.warning("⏳ 模型未训练")
    
    if st.session_state.prediction_results:
        st.success("✅ 预测已完成")
    else:
        st.info("💭 等待预测")